{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f26664cd-df92-4e2e-9701-93e4c72b0780",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !uv add loguru pydantic rich tenacity google-api-python-client beautifulsoup4\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2db3227-611a-4f41-9678-46a86568a58a",
   "metadata": {},
   "source": [
    "<img src=\"./docs/orchestration_workers1.jpg\" alt=\"Orchestration Workers 1\" width=\"1200\"/>\n",
    "<img src=\"./docs/orchestration_workers2.jpg\" alt=\"Orchestration Workers 2\" width=\"1200\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "549eec82-0f5b-400a-a1f5-992a42c970b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from typing import Any, Tuple\n",
    "\n",
    "import gradio as gr\n",
    "from agents import (\n",
    "    GuardrailFunctionOutput,\n",
    "    InputGuardrailTripwireTriggered,\n",
    "    Runner,\n",
    "    RunResult,\n",
    "    TResponseInputItem,\n",
    "    gen_trace_id,\n",
    "    trace,\n",
    ")\n",
    "from build_agents import planner_agent, routing_agent, search_agent, writer_agent\n",
    "from rich.console import Console\n",
    "from rich.markdown import Markdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88978524",
   "metadata": {},
   "outputs": [],
   "source": [
    "console = Console()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc6d2649-9206-4a4f-9d5e-ef241f8750bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResearchManager:\n",
    "    # The run method now assumes the query is pre-validated.\n",
    "    async def run(self, query: str) -> str:\n",
    "        \"\"\"\n",
    "        Runs the full research process on a pre-validated query.\n",
    "        Returns a markdown string of the report.\n",
    "        \"\"\"\n",
    "        # This trace block covers the actual research, not the validation.\n",
    "        with trace(workflow_name=\"Research trace\"):\n",
    "            # Step 1: Proceed directly to planning. The topic check is done outside.\n",
    "            plan: list[dict] = await self._plan_searches(query=query)\n",
    "            results = await self._perform_searches(plan)\n",
    "            report = await self._write_report(query, results)\n",
    "\n",
    "            return report.markdown_report\n",
    "\n",
    "    # This method will be called from the main loop and checks the validity of the query.\n",
    "    async def _topic_check(self, query: str) -> Tuple[bool, str]:\n",
    "        \"\"\"\n",
    "        Checks if the query is on-topic.\n",
    "        Returns (True, query) on success.\n",
    "        Returns (False, clarification_message) on failure.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            # The input must be in the chat history format for the guardrail.\n",
    "            formatted_input = [{\"role\": \"user\", \"content\": query}]\n",
    "            await Runner.run(starting_agent=routing_agent, input=formatted_input)\n",
    "\n",
    "            # If the guardrail does NOT trigger an exception, the topic is on.\n",
    "            return (True, query)\n",
    "\n",
    "        except InputGuardrailTripwireTriggered as e:\n",
    "            print(\"INFO | Guardrail | Topic check tripwire triggered.\")\n",
    "\n",
    "            try:\n",
    "                clarification_msg = e.guardrail_result.output.output_info[\n",
    "                    \"Found_Unexpected_Message\"\n",
    "                ].clarification_message\n",
    "            except (AttributeError, KeyError):\n",
    "                # If the expected structure isn't there, just use the default message.\n",
    "                print(\n",
    "                    \"INFO | Guardrail | Could not extract specific message, using default.\"\n",
    "                )\n",
    "                clarification_msg = \"Your request seems to be off-topic.\"\n",
    "\n",
    "            # print(f\"INFO | Guardrail | Clarification message: {clarification_msg}\")\n",
    "            return (False, clarification_msg)\n",
    "\n",
    "    async def _plan_searches(self, query: str):\n",
    "        result = await Runner.run(planner_agent, f\"Query: {query}\")\n",
    "        print(f\"INFO | Planning result: {result.final_output.searches}\")\n",
    "        return result.final_output.searches\n",
    "\n",
    "    async def _perform_searches(self, searches: list):\n",
    "        print(f\"INFO | Performing searches\")\n",
    "        tasks = [asyncio.create_task(self._search(item)) for item in searches]\n",
    "        summaries = await asyncio.gather(*tasks)\n",
    "        # print(f\"INFO | Search results: {summaries}\")\n",
    "        return summaries\n",
    "\n",
    "    async def _search(self, item: Any) -> str:\n",
    "        print(f\"INFO | - Searching for: {item.query}\")\n",
    "        res = await Runner.run(search_agent, item.query)\n",
    "        return res.final_output\n",
    "\n",
    "    async def _write_report(self, query: str, summaries: list[str]) -> Any:\n",
    "        print(f\"INFO | Writing report for query: {query}\")\n",
    "        joined = \"\\n\".join(summaries)\n",
    "        output = await Runner.run(\n",
    "            writer_agent, f\"Original query: {query}\\nResearch:\\n{joined}\"\n",
    "        )\n",
    "        return output.final_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2610463",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def main() -> None:\n",
    "    history: list[TResponseInputItem] = []\n",
    "    print(\"\\n\" + \"-\" * 98)\n",
    "    print(\n",
    "        \"Welcome to Coding Syllabus Generator! Ask me to generate syllabi about Python, JavaScript, or SQL.\"\n",
    "    )\n",
    "    print(\"Type 'quit' or 'exit' to end the conversation.\")\n",
    "    print(\"-\" * 98)\n",
    "\n",
    "    # Instantiate the manager once, outside the loop.\n",
    "    manager = ResearchManager()\n",
    "\n",
    "    while True:\n",
    "        message_str: str = input(\"\\nYou: \")\n",
    "        if not message_str or message_str.lower() in [\"quit\", \"exit\", \"thanks\", \"bye\"]:\n",
    "            print(\"\\nAssistant: Goodbye!\")\n",
    "            break\n",
    "\n",
    "        # The validation logic is here in the main loop.\n",
    "        final_result_string = \"\"\n",
    "        try:\n",
    "            # Step 1: Check if the topic is valid.\n",
    "            is_on_topic, check_result = await manager._topic_check(query=message_str)\n",
    "\n",
    "            # Step 2: If the topic is not valid, print the clarification and continue the loop.\n",
    "            if not is_on_topic:\n",
    "                # The 'check_result' is the clarification_message in this case.\n",
    "                final_result_string = check_result\n",
    "            else:\n",
    "                # Step 3: If the topic IS valid, run the full research process.\n",
    "                print(\"\\nINFO | Topic is valid. Starting research...\")\n",
    "                # The 'check_result' is the original query here, but we can just use message_str.\n",
    "                final_result_string = await manager.run(query=message_str)\n",
    "\n",
    "        except Exception as e:\n",
    "            # General error handling for the entire process\n",
    "            final_result_string = f\"An unexpected error occurred: {e}\"\n",
    "            # print(f\"ERROR | {final_result_string}\")\n",
    "\n",
    "        # This block now prints either the final report or the clarification message.\n",
    "        print(\"\\n\" + \"-\" * 10)\n",
    "        print(\"Assistant:\")\n",
    "        print(\"-\" * 10)\n",
    "        console.print(Markdown(markup=final_result_string))  \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b15abcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "await main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "agents",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
